#!/usr/bin/python
'''
Created by Colin Gillespie <c.gillespie@ncl.ac.uk>

Last Changed on: $Date$ 
Work in progress

'''

__revision__ = '$Revision$'

import sys, re
from sympy import diff

try:
    import psyco
    psyco.full()
except ImportError:
    pass

#LOOK AT CCODE in sympy!!!!!!!!!!!

#Non-standard modules
from mrange import cumRange
from VariableDict import Variables
from GetMomentEquation import getMomentEquation
from GetReactionInformation import GetStoichiometryDict
from GlobalVariables import globalVariables
import OutputFormat as OF 


class GSLInformation(object):
    '''Used as a storage container'''
    
    def __init__(self):
        self.pars = []
        self.species = []
        self.reactions_dict = {}
        self.max_values = [0]
        
    def putFileName(self, file_name):
        self.file_name = file_name
        self.fp = open(self.file_name, 'w')
    
    def getFile(self):
        return self.fp

    def closeFile(self):
        self.fp.close()
        return    
    
    def putStoicDict(self, gsd):
        self.gsd = gsd
     
    def getStoicDict(self):
        return self.gsd

    def putListOfPars(self, list_of_pars):
        self.pars = list_of_pars

    def getListOfPars(self):
        return self.pars
    
    

def getHeader(g):
    header = '''#include <stdio.h>
#include <math.h>    
#include <gsl/gsl_errno.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_odeiv.h>
gsl_matrix *rk_out;
#define MSET(rk_out,i,j,k) gsl_matrix_set(rk_out,i,j,k)\n
'''
    g.getFile().write(header)
    return True

def getParametersDict(GDD):#, list_of_pars):
    par_dict = {}
    for pa in GDD.pa_dict.keys():
        par_dict[pa] = str(GDD.pa_dict[pa]['value'])

    return par_dict


def createJacobianDict(gsl_info, value=None, reaction=None, jac_dict={}):
    if not value and not reaction:
        return jac_dict
    pos_i = OF.ConvertToGSL().order[value]

    for values in cumRange(gsl_info.max_values):
        pos_j = OF.ConvertToGSL().order[values]
        position = (pos_i, pos_j)
        moment = Variables().getMoment(values)
        jac_dict[position] = diff(reaction, moment)
        
def getjac(gsl_info):
    l = len(gsl_info.reactions_dict)
    expression = 'int jac(double t, const double EX[], double *dfdy,double dfdt[], void *params)\n{'
    expression += '\tgsl_matrix_view dfdy_mat = gsl_matrix_view_array (dfdy, %s, %s);\n' % (l, l)
    expression += '\tgsl_matrix *m = &dfdy_mat.matrix;\n'
    jac_dict = createJacobianDict(gsl_info)
    for value, key in jac_dict.items():
        expression += '\tgsl_matrix_set (m, %s, %s, %s);\n' % (value[0], value[1], str(key))

    for i in range(l):
        expression +='\tdfdt[%s]=0;\n' % str(i)
    
    expression += '\treturn GSL_SUCCESS;\n}'
    gsl_info.getFile().write(expression)
    return True

def getrk_main(gsl_info, initial_conditions_dict):
    l = len(gsl_info.reactions_dict)
    expression = '''int rk_main(int max_time)
    {
        const gsl_odeiv_step_type *T = gsl_odeiv_step_rk4;
        gsl_odeiv_step *stp = gsl_odeiv_step_alloc (T, %s);
        gsl_odeiv_control *c = gsl_odeiv_control_y_new (1e-6, 0.0);/*error allowances*/
        gsl_odeiv_evolve * e = gsl_odeiv_evolve_alloc (%s);
        gsl_odeiv_system sys = {func, 0, %s, 0};
  
        double t = 0.0;
        double h = 1e-6;
        int i,j;  
  
  ''' % (l, l, l)
    initial_conditions_list = l*[0]
    for value, key in initial_conditions_dict.items():
        index = OF.ConvertToGSL().order[value]
        initial_conditions_list[index] = float(key)


    ic = str(initial_conditions_list)[1:-1]
    expression += 'double EX[%s] = {%s};\n\n' % (l, ic)
    expression += '''for (i = 1; i <=(int) max_time; i++)
    {
     
      while (t < i)
	    {
	        gsl_odeiv_evolve_apply (e, c, stp, &sys, &t, i, &h, EX);
    	}
    for(j=0;j<3;j++) {
	    MSET(rk_out,j,i-1,EX[j]);
        }
    }


  gsl_odeiv_evolve_free (e);
  gsl_odeiv_control_free (c);
  gsl_odeiv_step_free (stp);
  return 0;
}\n
    '''
    gsl_info.getFile().write(expression)
    return True

def getMain(gsl_info):
    expression = '''int main()
    {
    int i,j;
    double max_time=100.0;
    rk_out=gsl_matrix_alloc(%s, (int) max_time);
  
    rk_main(max_time);
    for(i=0;i<max_time;i++){
        for(j=0;j<3;j++){
    '''% len(gsl_info.reactions_dict)
    expression += '''printf("%f ", gsl_matrix_get(rk_out,j,i));
        }
    printf("\\n"); 
    }
    return 0;
    }\n\n''' 
    gsl_info.getFile().write(expression)
    return True

def getfunc(gsl_info):
    expression = 'int func(double t, const double EX[], double dEX[], void *params)\n'
    expression += '{\n'
    m = gsl_info
   
    for value, key in m.getStoicDict().pa_dict.items():
        expression += '\tdouble %s=%s;\n' % (value, key['value'])
    p = re.compile('(EX\[\d*\])\*\*(\d*)', re.VERBOSE)

    for value, key in gsl_info.reactions_dict.items():
        reaction = p.sub(r'pow(\g<1>,\g<2>)', str(key))
        differ = Variables().getDifferential(value)
        expression += '\t%s =%s;\n' %(differ, reaction)
        createJacobianDict(gsl_info, value, key)
    expression += '\treturn GSL_SUCCESS;\n}\n\n'
    m.getFile().write(expression)
    return True

def getEqualsTo(gsl_info, specie_tup):
    equals_to = ''
    for i, value in enumerate(specie_tup):
        if value != 0:
            sp = gsl_info.getStoicDict().m.getSpecies(i)
            initial = sp.getInitialAmount()**value
            
            if len(equals_to):
                equals_to *=  initial
            else:
                equals_to = initial
    return str(equals_to)


def createGSLFile(sbml_file, output_file='tmp.c'):
    '''Creates the maple file with moment equations.
    
    Arguments:
        sbml_file - sbml file name
        output_file - outout file
    '''
    gsd = GetStoichiometryDict(sbml_file = sbml_file)
    gsl_info = GSLInformation()
    gsl_info.putFileName(output_file)
    gsl_info.putStoicDict(gsd)
    
    reaction_list = gsd.getReactionList()
    gsl_info.max_values = [globalVariables()['closure_order']] * len(gsd.empty_list)

    #Output the species
    initial_conditions_dict = {}

    
    for values in cumRange(gsl_info.max_values):
        initial_conditions_dict[values] = getEqualsTo(gsl_info, values)
        mom_eq = getMomentEquation(values, reaction_list)
        gsl_info.reactions_dict[values] = mom_eq
        if mom_eq != 0:
            gsl_info.species.append(values)
            #print '...Getting moment equation  dkappa_%s/dt' % str(values)
            #gsl_info.mom_equs.append(str(mom_eq))
        else:#mom_equ is constant
            if sum(values) > 1: #higher order moment
                raise 'Not sure if this has been test - sorry'
                #equals_to = getEqualsTo(m, values)
                #maple_sp = str(Variables().getMoment(tuple(values)))[:-3] + equals_to
                #gsl_info.pars.append(maple_sp)
            else:
                index = list(values).index(1)
                sp_id =  gsd.m.getSpecies(index).getId()
                print 'Note: Species %s is constant.' % sp_id
                print 'If you add - constant ="true" - to your SBML file then you will have fewer moment equations'
                
                gsl_info.species.append(values)


   
    getHeader(gsl_info)
    getfunc(gsl_info)
    getjac(gsl_info)
    getrk_main(gsl_info, initial_conditions_dict)
    getMain(gsl_info)
    return True

if __name__ == '__main__':
    globalVariables(output_format = 'GSL')
    createGSLFile('../xml_files/imdeath.xml')
    












